"""storage/index_manager.py
简单 B+Tree 索引管理器。
目前仅支持 INT 键、允许重复键 (链表存储多重 RID)。
功能：create / insert / search，自动处理叶子分裂及根提升。
后续可扩展 delete、范围查询等。
"""

from typing import List, Tuple
from .buffer import BufferManager
from .index_page import IndexPage, PAGE_TYPE_LEAF, PAGE_TYPE_INTERNAL, LEAF_ENTRY_SIZE, INTERNAL_ENTRY_SIZE, HEADER_SIZE, Page, PAGE_SIZE

RID = Tuple[int, int]

class IndexManager:
    def __init__(self, buffer_manager: BufferManager, root_page_id: int | None = None):
        self.buffer_manager = buffer_manager
        if root_page_id is None:
            # 创建根叶子页
            root_page_obj = buffer_manager.new_page()
            self.root_page_id = root_page_obj.page_id
            idx_page = IndexPage(root_page_obj)
            # 默认构造已经初始化为叶子
            buffer_manager.unpin_page(self.root_page_id, is_dirty=True)
        else:
            self.root_page_id = root_page_id

    # -------------- Public APIs --------------
    def search(self, key: int) -> List[RID]:
        """返回匹配 key 的 RID 列表，若不存在返回 []."""
        leaf_page = self._find_leaf(key)
        result = []
        for k, rid in leaf_page.iter_entries():
            if k == key:
                result.append(rid)
        self.buffer_manager.unpin_page(leaf_page.page.page_id)
        return result

    def insert(self, key: int, rid: RID):
        leaf_page = self._find_leaf(key)
        # 先尝试直接插入
        if leaf_page.insert_leaf_entry(key, rid):
            self.buffer_manager.unpin_page(leaf_page.page.page_id, is_dirty=True)
            return
        # 无空间，需要 split
        new_leaf, split_key = leaf_page.split_leaf(self.buffer_manager)
        # 确定插入目标
        target = new_leaf if key >= split_key else leaf_page
        target.insert_leaf_entry(key, rid)
        self.buffer_manager.unpin_page(leaf_page.page.page_id, is_dirty=True)
        self.buffer_manager.unpin_page(new_leaf.page.page_id, is_dirty=True)
        # 将 split_key 推到父节点
        self._insert_into_parent(leaf_page.page.page_id, split_key, new_leaf.page.page_id)

    # -------------- Internal helpers --------------
    def _find_leaf(self, key: int) -> IndexPage:
        page_obj = self.buffer_manager.get_page(self.root_page_id)
        idx_page = IndexPage(page_obj)
        while not idx_page.is_leaf:
            # 线性扫描 child 指针（尚未排序优化）
            child_page_id = None
            for k, child in idx_page.iter_entries():
                if key < k:
                    child_page_id = child
                    break
            if child_page_id is None:
                # key 大于所有 entry，取最后一个 child
                *_, last_child = list(idx_page.iter_entries())[-1]
                child_page_id = last_child
            self.buffer_manager.unpin_page(idx_page.page.page_id)
            page_obj = self.buffer_manager.get_page(child_page_id)
            idx_page = IndexPage(page_obj)
        return idx_page

    def _insert_into_parent(self, left_page_id: int, key: int, right_page_id: int):
        # 如果分裂的是根，需要新根
        if left_page_id == self.root_page_id:
            new_root_obj = self.buffer_manager.new_page()
            new_root_idx = IndexPage(new_root_obj)
            new_root_idx._init_page(PAGE_TYPE_INTERNAL, parent=0)
            new_root_idx.insert_internal_entry(key, left_page_id)
            new_root_idx.insert_internal_entry(2**31-1, right_page_id)  # sentinel

            # 更新子节点的 parent 指针
            left_pg = IndexPage(self.buffer_manager.get_page(left_page_id))
            right_pg = IndexPage(self.buffer_manager.get_page(right_page_id))
            left_pg.set_parent(new_root_obj.page_id)
            right_pg.set_parent(new_root_obj.page_id)
            self.buffer_manager.unpin_page(left_page_id, is_dirty=True)
            self.buffer_manager.unpin_page(right_page_id, is_dirty=True)

            self.root_page_id = new_root_obj.page_id
            self.buffer_manager.unpin_page(new_root_obj.page_id, is_dirty=True)
            return

        # 非根：通过父指针直接定位
        left_pg = IndexPage(self.buffer_manager.get_page(left_page_id))
        parent_id = left_pg.get_parent()
        self.buffer_manager.unpin_page(left_page_id)

        parent_obj = self.buffer_manager.get_page(parent_id)
        parent_idx = IndexPage(parent_obj)
        if parent_idx.insert_internal_entry(key, right_page_id):
            self.buffer_manager.unpin_page(parent_obj.page_id, is_dirty=True)
            return
        # 父节点也满 -> 递归分裂
        new_internal, split_key = parent_idx.split_internal(self.buffer_manager)
        target = new_internal if key >= split_key else parent_idx
        target.insert_internal_entry(key, right_page_id)
        self.buffer_manager.unpin_page(parent_obj.page_id, is_dirty=True)
        self.buffer_manager.unpin_page(new_internal.page.page_id, is_dirty=True)
        self._insert_into_parent(parent_id, split_key, new_internal.page.page_id)
